/*
 * Copyright 2010-2012 Susanta Tewari. <freecode4susant@users.sourceforge.net>
 *
 * This program is free software: you can redistribute it and/or modify
 * it under the terms of the GNU General Public License as published by
 * the Free Software Foundation, either version 3 of the License, or
 * (at your option) any later version.
 *
 * This program is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 * GNU General Public License for more details.
 *
 * You should have received a copy of the GNU General Public License
 * along with this program.  If not, see <http://www.gnu.org/licenses/>.
 */

package coalescent.phylogeny;

import coalescent.data.K69Data;
import coalescent.data.MutableK69Data;
import coalescent.model.K69;
import coalescent.phylogeny.GeneTree.MutationEdgeData;

import com.google.common.base.Function;
import com.google.common.base.Functions;

import commons.util.BinaryDataUtil;
import commons.util.ExceptionUtil;
import commons.util.graph.Edge;
import commons.util.graph.LeafPath;
import commons.util.graph.Node;

import java.util.*;

/**
 * Infinite Sites Data where the mutation history is maintained by an n-ary tree. Allele type is
 * {@code LeafPath} and the mutation type is {@code String}. Mutation information is encapsulated
 * in {@code MutationEdgeData}.
 *
 * @author Susanta Tewari
 */
public class GeneTree extends K69Data<LeafPath<String, MutationEdgeData>, String, GeneTree> {

    /**
     * this tree has only the distinct alleles
     */
    private final Node<String, MutationEdgeData> tree;

    /** Field description */
    private final int mutationCount;

    /** Field description */
    private final LeafPath<String, MutationEdgeData> wildAllele;

    /**
     * This is used internally to create a clone.
     *
     * @param model model
     * @param tree root node of the gene tree
     * @param alleleFreq allele frequency
     */
    private GeneTree(K69 model, Node<String, MutationEdgeData> tree,
                     Map<LeafPath<String, MutationEdgeData>, Integer> alleleFreq) {

        super(model, alleleFreq);

        this.mutationCount = mutationCount(tree);

        /**
         * Fix wild allele
         */


        // <editor-fold defaultstate="collapsed" desc="Fix wild allele">
        LeafPath<String, MutationEdgeData> wildPath = null;

        for (LeafPath<String, MutationEdgeData> path : alleleFreq.keySet()) {

            Node<String, MutationEdgeData> leaf      = path.getLeaf();
            Edge<String, MutationEdgeData> leaf_edge = leaf.getParent().getChildEdge(leaf);

            if (leaf.getParent().isRoot() && (leaf_edge.getData().mutationCount() == 0)) {

                wildPath = path;

                break;
            }
        }

        if (wildPath == null) {

            Collection<Node<String, MutationEdgeData>> nodes = new ArrayList<>(2);
            Node<String, MutationEdgeData> root              = new Node<>(null, tree.getData(), 1);
            Node<String, MutationEdgeData> leaf =
                root.addChild("wild", new MutationEdgeData(MutationEdgeData.NO_MUTATION_SYMBOL), 0);

            nodes.add(root);
            nodes.add(leaf);

            wildPath = new LeafPath<>(leaf);
        }

        this.wildAllele = wildPath;


        // </editor-fold>
        this.tree = tree;
    }

    /**
     * @param data infinite-sites data
     * @return a new instance of {@code GeneTree} based on the specified data
     */
    public static GeneTree getInstance(K69Data<String, Integer, ?> data) {


        // prepare data for gusfield`s algo
        Set<String> alleles              = data.getAlleles();
        Collection<byte[]> gusfield_data = new ArrayList<>(alleles.size());

        for (String allele : alleles) {

            gusfield_data.add(BinaryDataUtil.parseBinaryString(allele));
        }

        Node<String, String> ga_tree = BuildPPhylogenyAlgo.ALGO_GUSFIELD.buildTree(gusfield_data,
                                           MutationEdgeData.NO_MUTATION_SYMBOL,
                                           MutationEdgeData.MUTATION_DELIM);


        // transform tree using node and edge functions
        Function<String, String> nodeFunction           = Functions.identity();
        Function<String, MutationEdgeData> edgeFunction = new Function<String, MutationEdgeData>() {

            @Override
            public MutationEdgeData apply(String f) {
                return new MutationEdgeData(f);
            }
        };
        Node<String, MutationEdgeData> tree = Node.transform(ga_tree, nodeFunction, edgeFunction);


        // build allele freq from the tree
        List<Node<String, MutationEdgeData>> tree_leafs = tree.getLeafs();
        Map<LeafPath<String, MutationEdgeData>, Integer> alleleFreq =
            new HashMap<>(tree_leafs.size());

        for (Node<String, MutationEdgeData> leaf : tree_leafs) {

            LeafPath<String, MutationEdgeData> path = new LeafPath<>(leaf);
            String allele = BinaryDataUtil.buildBinaryString(alleleSequence(path));

            alleleFreq.put(path, data.getAlleleFrequency(allele));
        }

        return new GeneTree(data.getModel(), tree, alleleFreq);
    }

    /**
     * Method description
     *
     *
     * @return
     */
    public Node<String, MutationEdgeData> getRootNode() {
        return tree;
    }

    /**
     * Method description
     *
     *
     * @return
     */
    public Map<LeafPath<String, MutationEdgeData>, Integer> getNodeFreq() {

        Map<LeafPath<String, MutationEdgeData>, Integer> alleleFreq = getAlleleFreq();
        Map<LeafPath<String, MutationEdgeData>, Integer> result = new HashMap<>(alleleFreq.size());

        for (LeafPath<String, MutationEdgeData> allele : alleleFreq.keySet()) {

            result.put(allele, alleleFreq.get(allele));
        }

        return result;
    }

    /**
     * Method description
     *
     *
     * @param allele
     *
     * @return
     */
    private static byte[] alleleSequence(LeafPath<String, MutationEdgeData> allele) {

        byte[] result = new byte[allele.edgeLength()];
        int i         = 0;

        for (Edge<String, MutationEdgeData> edge : allele.edgeIterator()) {

            int mutationCount = edge.getData().mutationCount();

            if (mutationCount == 0) {
                result[i++] = 0;
            } else {

                for (int j = 0; j < mutationCount; j++) {

                    result[i++] = 1;
                }
            }
        }

        return result;
    }

    /**
     * Method description
     *
     *
     * @param allele
     *
     * @return
     */
    @Override
    public Set<String> getImmediateMutations(LeafPath<String, MutationEdgeData> allele) {

        checkMutationAllele(allele);

        Set<String> result = new HashSet<>(10);

        for (Edge<String, MutationEdgeData> edge : allele.edgeIterator()) {

            result.addAll(edge.getData().getMutations());
        }

        return result;
    }

    /**
     * Method description
     *
     *
     * @param allele
     * @param mutation
     *
     * @return
     */
    @Override
    public LeafPath<String, MutationEdgeData> mutateBackward(LeafPath<String,
                    MutationEdgeData> allele, String mutation) {

        checkMutationAllele(allele, mutation);

        LeafPath<String, MutationEdgeData> allele_copy = allele.copy();
        Edge<String, MutationEdgeData> mutation_edge   = allele.edgeIterator().iterator().next();
        MutationEdgeData new_edge_data = mutation_edge.getData().removeMutation(mutation);

        mutation_edge.setData(new_edge_data);

        return allele_copy;
    }

    /**
     * Method description
     *
     *
     * @return
     */
    @Override
    public LeafPath<String, MutationEdgeData> getWildAllele() {
        return wildAllele;
    }

    /**
     * Method description
     *
     *
     * @return
     */
    @Override
    public int getMutationCount() {
        return mutationCount;
    }

    /**
     * The number of mutations in the subtree under this node.
     *
     * @param node a node (internal or leaf) in this gene tree
     * @return number of mutations under this node.
     */
    private static int mutationCount(Node<String, MutationEdgeData> node) {

        int result = 0;

        for (Node<String, MutationEdgeData> child : node.getChildren()) {

            result += node.getChildEdge(child).getData().mutationCount();

            if (!child.isLeaf()) {
                result += mutationCount(child);
            }
        }

        return result;
    }

    /**
     * Method description
     *
     *
     * @return
     */
    @Override
    public MutableK69Data<LeafPath<String, MutationEdgeData>, String, GeneTree> getMutable() {
        throw new UnsupportedOperationException("Not supported yet.");
    }

    /**
     * Method description
     *
     *
     * @param data
     *
     * @return
     */
    @Override
    public GeneTree factory(GeneTree data) {
        throw new UnsupportedOperationException("Not supported yet.");
    }

    /**
     * Builds a path of mutation counts up along the edges all the way to the root, delimited by
     * {@link MutationEdgeData#MUTATION_DELIM MUTATION_DELIM}.
     *
     * @param node a node (internal or leaf) in this gene tree
     * @return upward path of mutation counts
     */
    protected String buildMutationCountPath(Node<String, MutationEdgeData> node) {

        StringBuilder stringBuilder = new StringBuilder();


        // parent-child init
        Node<String, MutationEdgeData> child      = node;
        Node<String, MutationEdgeData> the_parent = child.getParent();


        // move upward till the root
        while (!child.isRoot()) {


            // get the edge and its mutation count
            Edge<String, MutationEdgeData> edge = the_parent.getChildEdge(child);

            stringBuilder.append(edge.getData().mutationCount()).append(
                MutationEdgeData.MUTATION_DELIM);


            // parent-child swap
            child      = the_parent;
            the_parent = child.getParent();
        }


        // remove the last character
        if (stringBuilder.length() > 0) {
            stringBuilder.deleteCharAt(stringBuilder.length() - 1);
        }

        return stringBuilder.toString();
    }

    /**
     * Class description
     *
     *
     * @version        Enter version here..., 12/10/09
     * @author         Susanta Tewari
     */
    public static class MutationEdgeData {

        /** Field description */
        private static final String NO_MUTATION_SYMBOL = "0";

        /** Field description */
        public static final String MUTATION_DELIM = "-";

        /** Field description */
        private List<String> mutations;

        /**
         * Constructs ...
         *
         *
         * @param mutations
         */
        private MutationEdgeData(List<String> mutations) {
            this.mutations = mutations;
        }

        /**
         * Constructs ...
         *
         *
         * @param label
         */
        public MutationEdgeData(String label) {

            if (!label.equals(NO_MUTATION_SYMBOL)) {

                String[] mutation_temp = label.split(MUTATION_DELIM);

                mutations = new ArrayList<>(mutation_temp.length);

                mutations.addAll(Arrays.asList(mutation_temp));
            }
        }

        /**
         * Number of mutations on this edge
         *
         * @return number of mutations on this edge
         */
        public int mutationCount() {

            if (mutations == null) {
                return 0;
            }

            return mutations.size();
        }

        /**
         * Method description
         *
         *
         * @return
         */
        public List<String> getMutations() {

            if (mutations == null) {
                return Collections.emptyList();
            }

            return Collections.unmodifiableList(mutations);
        }

        /**
         * A factory method to create a new instance with the given mutation removed.
         *
         * @param mutation the mutation to be removed
         * @return a new instance of {@code MutationEdgeData} with the mutation removed
         */
        public MutationEdgeData removeMutation(String mutation) {

            if (!mutations.contains(mutation)) {
                ExceptionUtil.throwIllegalArgEx(getClass(), "Edge_Mutation_Absent", mutation);
            }


            // a copy for the new object
            Collection<String> new_mutations = new ArrayList<>(mutations);

            new_mutations.remove(mutation);

            return new MutationEdgeData(mutations);
        }
    }
}
